"""Module for sky image data model."""

import numpy as np
from scipy.interpolate import interp1d, RegularGridInterpolator, RectBivariateSpline
from astropy.io import fits
from astropy.coordinates import SkyCoord
import astropy.units as u


class SkyImage:
    """Data model for sky images."""
    
    def __init__(self, data=None, header=None, frequency=None, coordinates=None, nside=None):
        """
        Initialize SkyImage.
        
        Args:
            data: Image data array
            header: Image header information
            frequency: Frequency array corresponding to the data
            coordinates: Coordinate information for the image
            nside: NSIDE parameter for HEALPix maps
        """
        self.data = data
        self.header = header or {}
        self.frequency = frequency
        self.coordinates = coordinates
        self.nside = nside or 256  # Default to 256 for HEALPix maps
        self.is_healpix = nside is not None or (data is not None and data.ndim == 2 and data.shape[1] == 786432)
        
    def get_resolution(self):
        """Get image resolution."""
        if 'NAXIS1' in self.header and 'NAXIS2' in self.header:
            return (self.header['NAXIS1'], self.header['NAXIS2'])
        elif self.data is not None:
            if self.data.ndim == 2:
                return self.data.shape
            elif self.data.ndim == 3:
                return self.data.shape[1:]  # (time, height, width)
        return None
        
    def get_frequency(self):
        """Get image frequency."""
        return self.frequency
        
    def get_coordinates(self):
        """Get image coordinates."""
        return self.coordinates
        
    def interpolate_spatial(self, target_coordinates):
        """
        Interpolate sky image to target spatial coordinates.
        
        Args:
            target_coordinates: Target coordinate grid for interpolation
            
        Returns:
            Interpolated sky image data
        """
        if self.data is None:
            raise ValueError("Sky image data not available")
            
        # For HEALPix data, use healpy interpolation
        if self.is_healpix:
            try:
                import healpy as hp
                
                # Extract RA and Dec from the target coordinates
                ra = target_coordinates.ra.to(u.deg).value
                dec = target_coordinates.dec.to(u.deg).value
                
                # Convert to theta and phi (healpy uses radians)
                theta = np.pi/2 - np.deg2rad(dec)
                phi = np.deg2rad(ra)
                
                # For 2D HEALPix data (n_freq, n_pix)
                if self.data.ndim == 2:
                    # Create array to store interpolated values for each frequency
                    interpolated_data = np.zeros((self.data.shape[0], ra.size))
                    for i in range(self.data.shape[0]):
                        # Use healpy's get_interp_val for interpolation at each frequency
                        interpolated_data[i] = hp.get_interp_val(self.data[i], theta, phi)
                    return interpolated_data
                # For 1D HEALPix data (n_pix,)
                elif self.data.ndim == 1:
                    # Use healpy's get_interp_val for interpolation
                    interpolated_data = hp.get_interp_val(self.data, theta, phi)
                    return interpolated_data
                else:
                    raise NotImplementedError("HEALPix interpolation not implemented for this data shape")
            except ImportError:
                raise ImportError("healpy package is required for HEALPix interpolation")
        else:
            # For non-HEALPix data, use scipy interpolation
            # Assuming 2D data (height, width) for spatial interpolation
            if self.data.ndim >= 2:
                # Create coordinate grids for interpolation
                height, width = self.data.shape[-2:]
                x = np.arange(width)
                y = np.arange(height)
                
                # For 2D data
                if self.data.ndim == 2:
                    # Using RectBivariateSpline with explicit linear interpolation parameters
                    interpolator = RectBivariateSpline(y, x, self.data, kx=1, ky=1)
                    # Extract target coordinate values
                    x_target = target_coordinates[1]  # Assuming target_coordinates is (x, y)
                    y_target = target_coordinates[0]
                    # Use grid=False for proper non-grid interpolation
                    interpolated_data = interpolator(y_target, x_target, grid=False)
                # For 3D data (frequency, height, width)
                elif self.data.ndim == 3:
                    # Use RegularGridInterpolator for multi-dimensional interpolation
                    interpolator = RegularGridInterpolator((np.arange(self.data.shape[0]), y, x), 
                                                         self.data, method='linear')
                    # Assuming target_coordinates is properly shaped for 3D interpolation
                    interpolated_data = interpolator(target_coordinates)
                else:
                    raise NotImplementedError("Interpolation for this data dimensionality not implemented")
                    
                return interpolated_data
            else:
                raise ValueError("Data must be at least 2D for spatial interpolation")
        
    def interpolate_frequency(self, target_frequencies):
        """
        Interpolate sky image to target frequencies.
        
        Args:
            target_frequencies: Target frequencies for interpolation
            
        Returns:
            Interpolated sky image data
        """
        if self.data is None:
            raise ValueError("Sky image data not available")
            
        # For HEALPix data with frequency information
        if self.is_healpix and self.frequency is not None:
            # Use 1D interpolation along the frequency axis
            # Assuming self.data has shape (n_freq, n_pix)
            interpolator = interp1d(self.frequency[:-1], self.data, axis=0, kind='linear', 
                                  fill_value='extrapolate')
            interpolated_data = interpolator(target_frequencies)
            return interpolated_data
        # For regular image data with frequency information
        elif self.frequency is not None and self.data.ndim >= 3:
            # Assuming data has shape (n_freq, height, width) or more dimensions
            # Use interpolation along the first axis (frequency)
            interpolator = interp1d(self.frequency[:-1], self.data, axis=0, kind='linear',
                                  fill_value='extrapolate')
            interpolated_data = interpolator(target_frequencies)
            return interpolated_data
        else:
            raise ValueError("Frequency information not available or data format not supported")
        
    def get_intensity_at_coordinate(self, coordinate, frequency=None):
        """
        Get intensity at a specific coordinate using interpolation.
        
        Args:
            coordinate: SkyCoord object representing the coordinate to extract intensity from
            frequency: Target frequency for interpolation (optional)
            
        Returns:
            Intensity at the specified coordinate
        """
        if self.data is None:
            raise ValueError("Sky image data not available")
            
        # Check if we have HEALPix data
        if self.is_healpix:
            # For HEALPix data, directly interpolate at the coordinate
            interpolated_data = self.interpolate_spatial(coordinate)
            
            # If frequency is specified, interpolate along frequency axis
            if frequency is not None and self.frequency is not None:
                # Interpolate frequency dimension
                freq_interpolator = interp1d(self.frequency[:-1], interpolated_data, axis=0, kind='linear',
                                           fill_value='extrapolate')
                final_data = freq_interpolator(frequency)
                return final_data
            else:
                return interpolated_data
        else:
            # For non-HEALPix data, other interpolation methods could be used
            # This is a simplified implementation
            # In practice, you would need to map the coordinate to pixel coordinates
            # and then interpolate
            interpolated_data = self.interpolate_spatial(coordinate)
            
            # If frequency is specified, interpolate along frequency axis
            if frequency is not None and self.frequency is not None:
                # Interpolate frequency dimension
                freq_interpolator = interp1d(self.frequency[:-1], interpolated_data, axis=0, kind='linear',
                                           fill_value='extrapolate')
                final_data = freq_interpolator(frequency)
                return final_data
            else:
                return interpolated_data